package ssau.knyazev.fft.modules.oob;

import ssau.knyazev.core.B2Complex;
import ssau.knyazev.core.Complex;
import ssau.knyazev.core.CComplex;
import ssau.knyazev.modules.Generator;

public class B2FFTModule {

	protected B2FFTModule() {	}

	public static Complex[][] fft2D(float[][] src) {
		B2Complex[][] arr = transform(src);
		return fft2D(arr);
	}

	public static Complex[][] fft2D(Complex[][] src) {
		B2Complex[][] arr = transform(src);
		return fft2D(arr);
	}

	public static Complex[][] fft2D(B2Complex[][] src) {
		inverse(src);
		
		int x = src.length;
		int N = src.length;
		int size = 0;
		while (x > 1) {
			size++;
			x = x >> 1;
		}
		B2Complex[] wi_list = new B2Complex[src.length];
		B2Complex[] wj_list = new B2Complex[src.length];
		float cos, sin;
		for (int i = 0; i < src.length; i++) {
			double power = -2*Math.PI*i/N;
			cos = (float)Math.cos(power)/2;
			sin = (float)Math.sin(power)/2;
			wi_list[i] = new B2Complex(cos, cos, sin, sin);
			wj_list[i] = new B2Complex(cos, cos, -sin, sin);
		}
		for (int q = 0; q < size; q++) {
			int step = 1 << q;
			int step2 = step << 1;
			for (int down = 0; down < N - step; down++) {
				for (int right = 0; right < N - step; right++) {
					boolean flag = (down/step) % 2 + (right/step) % 2 > 0 ? false : true;
					if (flag) {
						B2Complex t1 = wi_list[(down % step2)*(N/step2)];
						B2Complex t2 = wj_list[(right % step2)*(N/step2)];
						fft2d2x2U0U2(src[down][right], src[down][right + step], src[down + step][right], src[down + step][right + step], t1, t2);
						fft2d2x2U1U3(src[down][right], src[down][right + step], src[down + step][right], src[down + step][right + step], t1, t2);
					}
				}
			}
		}
		return null;//convertToComplex(src);
	}
	


	protected static void fft2d2x2U1U3(B2Complex a00, B2Complex a01, B2Complex a10, B2Complex a11, B2Complex wi, B2Complex wj) {
		B2Complex newa00 = new B2Complex();
		B2Complex newa01 = new B2Complex();
		B2Complex newa10 = new B2Complex();
		B2Complex newa11 = new B2Complex();

		newa00 = B2Complex.sumU1U3(newa00, a00);
		newa10 = B2Complex.sumU1U3(newa10, a00);
		newa01 = B2Complex.sumU1U3(newa01, a00);
		newa11 = B2Complex.sumU1U3(newa11, a00);
		
		B2Complex t = B2Complex.multyU1U3(wi, a10);
		newa00 = B2Complex.sumU1U3(newa00, t);
		newa10 = B2Complex.subU1U3(newa10, t);
		newa01 = B2Complex.sumU1U3(newa01, t);
		newa11 = B2Complex.subU1U3(newa11, t);
		
		t = B2Complex.multyU1U3(wj, a01);
		newa00 = B2Complex.sumU1U3(newa00, t);
		newa10 = B2Complex.sumU1U3(newa10, t);
		newa01 = B2Complex.subU1U3(newa01, t);
		newa11 = B2Complex.subU1U3(newa11, t);
		
		t = B2Complex.multyU1U3(wj, a11);
		t = B2Complex.multyU1U3(wi, t);
		newa00 = B2Complex.sumU1U3(newa00, t);
		newa10 = B2Complex.subU1U3(newa10, t);
		newa01 = B2Complex.subU1U3(newa01, t);
		newa11 = B2Complex.sumU1U3(newa11, t);

		copyValuesU1U3(a00, newa00);
		copyValuesU1U3(a01, newa01);
		copyValuesU1U3(a10, newa10);
		copyValuesU1U3(a11, newa11);
	}
	
	protected static void fft2d2x2U0U2(B2Complex a00, B2Complex a01, B2Complex a10, B2Complex a11, B2Complex wi, B2Complex wj) {
		B2Complex newa00 = new B2Complex();
		B2Complex newa01 = new B2Complex();
		B2Complex newa10 = new B2Complex();
		B2Complex newa11 = new B2Complex();

		newa00 = B2Complex.sumU0U2(newa00, a00);
		newa10 = B2Complex.sumU0U2(newa10, a00);
		newa01 = B2Complex.sumU0U2(newa01, a00);
		newa11 = B2Complex.sumU0U2(newa11, a00);
		
		B2Complex t = B2Complex.multyU0U2(wi, a10);
		newa00 = B2Complex.sumU0U2(newa00, t);
		newa10 = B2Complex.subU0U2(newa10, t);
		newa01 = B2Complex.sumU0U2(newa01, t);
		newa11 = B2Complex.subU0U2(newa11, t);
		
		t = B2Complex.multyU0U2(wj, a01);
		newa00 = B2Complex.sumU0U2(newa00, t);
		newa10 = B2Complex.sumU0U2(newa10, t);
		newa01 = B2Complex.subU0U2(newa01, t);
		newa11 = B2Complex.subU0U2(newa11, t);
		
		t = B2Complex.multyU0U2(wj, a11);
		t = B2Complex.multyU0U2(wi, t);
		newa00 = B2Complex.sumU0U2(newa00, t);
		newa10 = B2Complex.subU0U2(newa10, t);
		newa01 = B2Complex.subU0U2(newa01, t);
		newa11 = B2Complex.sumU0U2(newa11, t);

		copyValuesU0U2(a00, newa00);
		copyValuesU0U2(a01, newa01);
		copyValuesU0U2(a10, newa10);
		copyValuesU0U2(a11, newa11);
	}
	
	protected static B2Complex[][] union(B2Complex[][] arr1, B2Complex[][] arr2){
		B2Complex[][] res = new B2Complex[arr1.length][arr1[0].length];
		for (int i = 0; i < arr1.length; i++){
			for (int j = 0; j < arr1[i].length; j++){
				res[i][j] = B2Complex.sum(arr1[i][j], arr2[i][j]);
			}
		}
		return res;
	}
	
	protected static void copyValues(B2Complex a, B2Complex b) {
		a.setU0(b.getU0());
		a.setU1(b.getU1());
		a.setU2(b.getU2());
		a.setU3(b.getU3());
	}
	
	protected static void copyValuesU0U2(B2Complex a, B2Complex b) {
		a.setU0(b.getU0());
		a.setU2(b.getU2());
	}
	
	protected static void copyValuesU1U3(B2Complex a, B2Complex b) {
		a.setU1(b.getU1());
		a.setU3(b.getU3());
	}

	protected static void inverse(B2Complex[] arr) {
		int dl = arr.length / 2;
		int st = arr.length - 1;
		int j = 0;
		for (int i = 0; i < st; i++) {
			if (i < j) {
				B2Complex s0 = arr[i];
				arr[i] = arr[j];
				arr[j] = s0;
			}
			int k = dl;
			while (k <= j) {
				j = j - k;
				k = k >> 1;
			}
			j = j + k;
		}
	}

	protected static void inverse(B2Complex[][] src) {
		int N = src.length;
		for (int i = 0; i < N; i++) {
			inverse(src[i]);
		}
		B2Complex[] arr = new B2Complex[N];
		for (int i = 0; i < N; i++) {
			for (int j = 0; j < N; j++) {
				arr[j] = src[j][i];
			}
			inverse(arr);
			for (int j = 0; j < N; j++) {
				src[j][i] = arr[j];
			}
		}
	}
	
	protected static B2Complex[][] transform(float[][] in) {
		B2Complex[][] res = new B2Complex[in.length][in.length];
		for (int i = 0; i < in.length; i++)
			for (int j = 0; j < in.length; j++) {
				res[i][j] = new B2Complex(in[i][j]);
			}
		return res;
	}

	protected static B2Complex[][] transform(Complex[][] in) {
		B2Complex[][] res = new B2Complex[in.length][in.length];
		for (int i = 0; i < in.length; i++){
			for (int j = 0; j < in.length; j++) {
				res[i][j] = new B2Complex(new CComplex(in[i][j].getReal(), 0, 0, 0));
			}
		}
		return res;
	}
	

	protected static CComplex[][] convertToCComplex(B2Complex[][] arr) {
		CComplex[][] res = new CComplex[arr.length][arr.length];
		for (int i = 0; i < res.length; i++)
			for (int j = 0; j < res.length; j++)
				res[i][j] = arr[i][j].getCComplex();
		return res;
	}
	
	protected static Complex[][] convertToComplex(CComplex[][] arr) {
		Complex[][] res = new Complex[arr.length][arr.length];
		for (int i = 0; i < res.length; i++){
			for (int j = 0; j < res.length; j++){
				res[i][j] = arr[i][j].getComplex();
			}
		}
		return res;
	}
	
	protected static Complex[][] convertToComplex(B2Complex[][] arr) {
		Complex[][] res = new Complex[arr.length][arr.length];
		for (int i = 0; i < res.length; i++){
			for (int j = 0; j < res.length; j++){
				res[i][j] = arr[i][j].getComplex();
			}
		}
		return res;
	}
	
	protected static B2Complex[][] copyMatrix(B2Complex[][] src){
		B2Complex[][] res = new B2Complex[src.length][src[0].length];
		for (int i = 0; i < res.length; i++){
			for (int j = 0; j < res.length; j++){
				res[i][j] = src[i][j].clone();
			}
		}
		return res;
	}
	
	public static void main(String[] args){
		float[][] arr = Generator.generateSimpleMatrix(4096);
		//Printer.printVector(arr);
		long sum = 0;
		for (int i = 0; i < 2; i++){
			System.out.println(i);
			long time = System.nanoTime();
			Complex[][] res = fft2D(arr);
			time = System.nanoTime() - time;
			sum += time;
		}
		System.out.println("RESULT");
		//Printer.printVector(res);
		System.out.println("\n" + sum/2);
		
	}

}
